{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "dcd81ae0",
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "cannot import name 'HybridSearcher' from 'pyserini.search' (/home/xueguang/anaconda3/envs/binder/lib/python3.8/site-packages/pyserini/search/__init__.py)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "Input \u001b[0;32mIn [38]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpyserini\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msearch\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m FaissSearcher, LuceneSearcher, HybridSearcher\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpyserini\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msearch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfaiss\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AutoQueryEncoder\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpyserini\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msearch\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m get_topics, get_qrels\n",
      "\u001b[0;31mImportError\u001b[0m: cannot import name 'HybridSearcher' from 'pyserini.search' (/home/xueguang/anaconda3/envs/binder/lib/python3.8/site-packages/pyserini/search/__init__.py)"
     ]
    }
   ],
   "source": [
    "from pyserini.search import FaissSearcher, LuceneSearcher\n",
    "from pyserini.\n",
    "from pyserini.search.faiss import AutoQueryEncoder\n",
    "from pyserini.search import get_topics, get_qrels\n",
    "import json\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "33ae18b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus = LuceneSearcher('index_relation_fb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6d10e10e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# BM25 search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "e1fe2437",
   "metadata": {},
   "outputs": [],
   "source": [
    "bm25_searcher = LuceneSearcher('index_relation_fb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "068292f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': '18202', 'contents': 'measurement unit surface tension unit tension in newtons per meter', 'rel_ori': 'measurement_unit.surface_tension_unit.tension_in_newtons_per_meter'}\n",
      "{'id': '18203', 'contents': 'measurement unit surface tension unit measurement system', 'rel_ori': 'measurement_unit.surface_tension_unit.measurement_system'}\n",
      "{'id': '20530', 'contents': 'measurement unit measurement system surface tension units', 'rel_ori': 'measurement_unit.measurement_system.surface_tension_units'}\n",
      "{'id': '18525', 'contents': 'measurement unit unit of surface density measurement system', 'rel_ori': 'measurement_unit.unit_of_surface_density.measurement_system'}\n",
      "{'id': '20509', 'contents': 'measurement unit measurement system surface density units', 'rel_ori': 'measurement_unit.measurement_system.surface_density_units'}\n",
      "{'id': '17153', 'contents': 'measurement unit unit of volumetric flow rate rate in cubic metres per second', 'rel_ori': 'measurement_unit.unit_of_volumetric_flow_rate.rate_in_cubic_metres_per_second'}\n",
      "{'id': '18526', 'contents': 'measurement unit unit of surface density density in kilograms per square meter', 'rel_ori': 'measurement_unit.unit_of_surface_density.density_in_kilograms_per_square_meter'}\n",
      "{'id': '17162', 'contents': 'measurement unit unit of force force in newtons', 'rel_ori': 'measurement_unit.unit_of_force.force_in_newtons'}\n",
      "{'id': '11299', 'contents': 'measurement unit force moment unit moment in newton meters', 'rel_ori': 'measurement_unit.force_moment_unit.moment_in_newton_meters'}\n",
      "{'id': '11246', 'contents': 'measurement unit unit of conductivity measurement system', 'rel_ori': 'measurement_unit.unit_of_conductivity.measurement_system'}\n",
      "{'id': '11518', 'contents': 'measurement unit acceleration unit measurement system', 'rel_ori': 'measurement_unit.acceleration_unit.measurement_system'}\n",
      "{'id': '11675', 'contents': 'measurement unit distance unit measurement system', 'rel_ori': 'measurement_unit.distance_unit.measurement_system'}\n",
      "{'id': '11719', 'contents': 'measurement unit permeability unit measurement system', 'rel_ori': 'measurement_unit.permeability_unit.measurement_system'}\n",
      "{'id': '11789', 'contents': 'measurement unit substance unit measurement system', 'rel_ori': 'measurement_unit.substance_unit.measurement_system'}\n",
      "{'id': '13070', 'contents': 'measurement unit energy unit measurement system', 'rel_ori': 'measurement_unit.energy_unit.measurement_system'}\n",
      "{'id': '13221', 'contents': 'measurement unit unit of density measurement system', 'rel_ori': 'measurement_unit.unit_of_density.measurement_system'}\n",
      "{'id': '13256', 'contents': 'measurement unit wavenumber unit measurement system', 'rel_ori': 'measurement_unit.wavenumber_unit.measurement_system'}\n",
      "{'id': '13360', 'contents': 'measurement unit measurement system luminance units', 'rel_ori': 'measurement_unit.measurement_system.luminance_units'}\n",
      "{'id': '13361', 'contents': 'measurement unit measurement system speed units', 'rel_ori': 'measurement_unit.measurement_system.speed_units'}\n",
      "{'id': '13362', 'contents': 'measurement unit measurement system density units', 'rel_ori': 'measurement_unit.measurement_system.density_units'}\n"
     ]
    }
   ],
   "source": [
    "hits = bm25_searcher.search('newton per metre is the unit of measurement for surface tension in which system of measurement?', k=20)\n",
    "for hit in hits:\n",
    "    print(json.loads(corpus.doc(str(hit.docid)).raw()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "4f9b1f66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Contriever Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "98d42f86",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_encoder = AutoQueryEncoder(encoder_dir='facebook/contriever', pooling='mean')\n",
    "contriever_searcher = FaissSearcher('freebase_contriever_index', query_encoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "6cd63532",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': '18202', 'contents': 'measurement unit surface tension unit tension in newtons per meter', 'rel_ori': 'measurement_unit.surface_tension_unit.tension_in_newtons_per_meter'}\n",
      "{'id': '20530', 'contents': 'measurement unit measurement system surface tension units', 'rel_ori': 'measurement_unit.measurement_system.surface_tension_units'}\n",
      "{'id': '18525', 'contents': 'measurement unit unit of surface density measurement system', 'rel_ori': 'measurement_unit.unit_of_surface_density.measurement_system'}\n",
      "{'id': '17162', 'contents': 'measurement unit unit of force force in newtons', 'rel_ori': 'measurement_unit.unit_of_force.force_in_newtons'}\n",
      "{'id': '18526', 'contents': 'measurement unit unit of surface density density in kilograms per square meter', 'rel_ori': 'measurement_unit.unit_of_surface_density.density_in_kilograms_per_square_meter'}\n",
      "{'id': '18203', 'contents': 'measurement unit surface tension unit measurement system', 'rel_ori': 'measurement_unit.surface_tension_unit.measurement_system'}\n",
      "{'id': '17153', 'contents': 'measurement unit unit of volumetric flow rate rate in cubic metres per second', 'rel_ori': 'measurement_unit.unit_of_volumetric_flow_rate.rate_in_cubic_metres_per_second'}\n",
      "{'id': '19608', 'contents': 'measurement unit unit of force measurement system', 'rel_ori': 'measurement_unit.unit_of_force.measurement_system'}\n",
      "{'id': '11299', 'contents': 'measurement unit force moment unit moment in newton meters', 'rel_ori': 'measurement_unit.force_moment_unit.moment_in_newton_meters'}\n",
      "{'id': '2433', 'contents': 'measurement unit money value amount', 'rel_ori': 'measurement_unit.money_value.amount'}\n",
      "{'id': '15451', 'contents': 'measurement unit angular velocity unit frequency in radians per second', 'rel_ori': 'measurement_unit.angular_velocity_unit.frequency_in_radians_per_second'}\n",
      "{'id': '6826', 'contents': 'measurement unit electric field strength unit strength in volts per meter', 'rel_ori': 'measurement_unit.electric_field_strength_unit.strength_in_volts_per_meter'}\n",
      "{'id': '7526', 'contents': 'measurement unit pressure unit pressure in pascals', 'rel_ori': 'measurement_unit.pressure_unit.pressure_in_pascals'}\n",
      "{'id': '18470', 'contents': 'measurement unit measurement system force moment units', 'rel_ori': 'measurement_unit.measurement_system.force_moment_units'}\n",
      "{'id': '20812', 'contents': 'measurement unit volumetric heat capacity unit capacity in joules per cubic meter per kelvin', 'rel_ori': 'measurement_unit.volumetric_heat_capacity_unit.capacity_in_joules_per_cubic_meter_per_kelvin'}\n",
      "{'id': '20509', 'contents': 'measurement unit measurement system surface density units', 'rel_ori': 'measurement_unit.measurement_system.surface_density_units'}\n",
      "{'id': '18467', 'contents': 'base engineering structural steel designation mass per metre', 'rel_ori': 'base.engineering.structural_steel_designation.mass_per_metre'}\n",
      "{'id': '8926', 'contents': 'measurement unit mass unit weightmass in kilograms', 'rel_ori': 'measurement_unit.mass_unit.weightmass_in_kilograms'}\n",
      "{'id': '14807', 'contents': 'measurement unit magnetic field strength unit strength in amperes per meter', 'rel_ori': 'measurement_unit.magnetic_field_strength_unit.strength_in_amperes_per_meter'}\n",
      "{'id': '13221', 'contents': 'measurement unit unit of density measurement system', 'rel_ori': 'measurement_unit.unit_of_density.measurement_system'}\n"
     ]
    }
   ],
   "source": [
    "hits = contriever_searcher.search('newton per metre is the unit of measurement for surface tension in which system of measurement?', k=20)\n",
    "for hit in hits:\n",
    "    print(json.loads(corpus.doc(str(hit.docid)).raw()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "88b96974",
   "metadata": {},
   "outputs": [],
   "source": [
    "# hybrid searcher that combine BM25 and Contriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "3b55ce7a",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'HybridSearcher' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Input \u001b[0;32mIn [37]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m hsearcher \u001b[38;5;241m=\u001b[39m \u001b[43mHybridSearcher\u001b[49m(contriever_searcher, bm25_searcher)\n\u001b[1;32m      2\u001b[0m hits \u001b[38;5;241m=\u001b[39m hsearcher\u001b[38;5;241m.\u001b[39msearch(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mnewton per metre is the unit of measurement for surface tension in which system of measurement?\u001b[39m\u001b[38;5;124m'\u001b[39m, k\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m20\u001b[39m)\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m hit \u001b[38;5;129;01min\u001b[39;00m hits:\n",
      "\u001b[0;31mNameError\u001b[0m: name 'HybridSearcher' is not defined"
     ]
    }
   ],
   "source": [
    "hsearcher = HybridSearcher(contriever_searcher, bm25_searcher)\n",
    "hits = hsearcher.search('newton per metre is the unit of measurement for surface tension in which system of measurement?', k=20)\n",
    "for hit in hits:\n",
    "    print(json.loads(corpus.doc(str(hit.docid)).raw()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dbf65c1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
